import jieba
import numpy as np
from wordcloud import WordCloud, ImageColorGenerator
from PIL import Image
import matplotlib.pyplot as plt
import re

# 对评论进行预处理
def preprocess(comment):
    comment = comment.lower()
    comment = re.sub(r'[^\u4e00-\u9fa5]', '', comment)
    return comment

# 提取关键词
def extract_keywords(comment):
    comment = preprocess(comment)
    words = jieba.cut(comment, cut_all=False)  # 分词
    keywords = []

    for word in words:
        if len(word) > 1 and word not in ['的', '和', '了']:
            keywords.append(word)

    return keywords

# 创建词云图片
def create_wordcloud(comments):
    # 合并所有评论
    all_comments = ' '.join(comments)

    # 提取关键词
    keywords = extract_keywords(all_comments)
    # 返回词云图片
    # get_wordcloud(keywords)
    # 返回定制图案的词云图片
    get_wordcloud_with_image(keywords)



def get_wordcloud(comments):
    # 创建词云
    wordcloud = WordCloud(
        font_path='msyh.ttc',  # 设置中文字体
        background_color='white',  # 设置背景颜色
        max_words=100,  # 设置最大显示的词数
        width=800,  # 设置图片宽度
        height=400,  # 设置图片高度
        stopwords=['的', '和', '了']  # 设置停用词
    ).generate(' '.join(comments))

    # 显示词云图片
    plt.imshow(wordcloud, interpolation='bilinear')
    plt.axis('off')
    plt.show()


def get_wordcloud_with_image(comments):
    # 读取图片文件
    image = np.array(Image.open("mask/img_1.png"))

    # 创建词云
    wordcloud = WordCloud(
        font_path='msyh.ttc',  # 设置中文字体
        background_color='white',  # 设置背景颜色
        max_words=100,  # 设置最大显示的词数
        width=800,  # 设置图片宽度
        height=400,  # 设置图片高度
        stopwords=['的', '和', '了'],  # 设置停用词
        mask=image
    ).generate(' '.join(comments))

    # 创建ImageColorGenerator对象，用于从图片中提取颜色
    image_colors = ImageColorGenerator(image)

    # 显示词云
    plt.imshow(wordcloud.recolor(color_func=image_colors), interpolation="bilinear")
    plt.axis("off")
    plt.show()


comments = [
    "这个产品很好用，性价比很高。",
    "非常满意，物有所值。",
    "很不错，值得购买。",
    "该商品非常棒，我很喜欢。质量非常好，价格也很合理。物流也很给力，非常满意！",
    "发货很快，颜色、质地很好，就是有些薄。总体来说不错，继续支持!",
    "老板性格好，宝贝也好，戴上去也很舒服，超赞.",
    "掌柜人不错。鞋子很好服务很热情。"
]

create_wordcloud(comments)
